import os
import json
from PyCmpltrtok.common import get_dir_name_ext, sep

__None__ = '__None__'


def map2set_2_map2list(xdict):
    return {k: sorted(v) for k, v in xdict.items()}


def check_spo_jsonl(xpath, limit=0):
    xs_type_set = set()
    xp_set = set()
    xo_type_set = set()

    xsub_type2sub_set_map = dict()
    xobj_type2obj_value_set_map = dict()

    xs_none_cnt = 0
    xs_type_none_cnt = 0
    xp_none_cnt = 0
    xo_none_cnt = 0
    xo_type_none_cnt = 0

    xother_key_of_object = set()
    xother_key_of_object_type = set()

    with open(xpath, 'r', encoding='utf8') as fi:
        cnt = -1
        while True:
            # 读取一行
            xline = fi.readline()
            if not xline:
                break
            if '\r\n' == xline[-2:]:
                xline = xline[:-2]
            elif '\r' == xline[-1:]:
                xline = xline[:-1]
            xline = xline.strip()
            if not xline:
                continue

            # 人为限定读取的行数
            cnt += 1
            if limit and cnt >= limit:
                break

            # 解析
            xdict = json.loads(xline)
            spo_list = xdict['spo_list']
            for xspo in spo_list:
                # S
                xs = xspo.get('subject', None)
                if xs is None:
                    xs_none_cnt += 1
                xs_type = xspo.get('subject_type', None)
                if xs_type is None:
                    xs_type_none_cnt += 1
                    xs_type = __None__

                # 统计S
                if xs_type is not None:
                    xs_type_set.add(xs_type)
                    if xs is not None:
                        xset = xsub_type2sub_set_map.get(xs_type, None)
                        if xset is None:
                            xset = set()
                        xset.add(xs)
                        xsub_type2sub_set_map[xs_type] = xset

                # P
                xp = xspo.get('predicate', None)
                if xp is None:
                    xp_none_cnt += 1

                # 统计P
                if xp is not None:
                    xp_set.add(xp)

                # object@value
                xo = xspo.get('object', None)
                if xo is None:
                    xo_none_cnt += 1
                else:
                    xkeys = set(xo.keys())
                    xkeys.discard('@value')
                    if xkeys:
                        xother_key_of_object = xother_key_of_object.union(xkeys)
                    xo = xo.get('@value', None)
                    if xo is None:
                        xo_none_cnt += 1

                # object_type@value
                xo_type = xspo.get('object_type', None)
                if xo_type is None:
                    xo_type_none_cnt += 1
                else:
                    xkeys = set(xo_type.keys())
                    xkeys.discard('@value')
                    if xkeys:
                        xother_key_of_object_type = xother_key_of_object_type.union(xkeys)
                    xo_type = xo_type.get('@value', None)
                    if xo_type is None:
                        xo_type_none_cnt += 1
                        xo_type = __None__

                # 统计object
                if xo_type is not None:
                    xo_type_set.add(xo_type)
                    if xo is not None:
                        xset = xobj_type2obj_value_set_map.get(xo_type, None)
                        if xset is None:
                            xset = set()
                        xset.add(xo)
                        xobj_type2obj_value_set_map[xo_type] = xset
    return {
        'sub_type': sorted(xs_type_set),
        'sub_type2sub': map2set_2_map2list(xsub_type2sub_set_map),
        'pred': sorted(xp_set),
        'obj_type': sorted(xo_type_set),
        'obj_type2obj': map2set_2_map2list(xobj_type2obj_value_set_map),
        'sub_none_cnt': xs_none_cnt,
        'sub_type_none_cnt': xs_type_none_cnt,
        'pred_none_cnt': xp_none_cnt,
        'obj_none_cnt': xo_none_cnt,
        'obj_type_none_cnt': xo_type_none_cnt,
        'other_key_of_object': sorted(xother_key_of_object),
        'other_key_of_object_type': sorted(xother_key_of_object_type),
    }


if '__main__' == __name__:

    def check(spo_jsonl_path, limit=20):
        sep(f'Read {spo_jsonl_path} with limit {limit}')
        if limit:
            limit_str = f'.limit-{limit}'
        else:
            limit_str = ''
        xdict = check_spo_jsonl(spo_jsonl_path, limit)
        print(xdict)
        xdir, xbase, xext = get_dir_name_ext(spo_jsonl_path)
        spo_check_path = os.path.join(xdir, xbase + f'.check{limit_str}.json')
        sep('Dump to ' + spo_check_path)
        with open(spo_check_path, 'w', encoding='utf8') as fo:
            json.dump(xdict, fo)
        sep('Over')

    def _main():
        xlist = [
            r'D:\_dell7590_root\local\LNP_datasets\med\CBLUE\CMeIE-V2\CMeIE-V2_train.jsonl',
            r'D:\_dell7590_root\local\LNP_datasets\med\CBLUE\CMeIE-V2\CMeIE-V2_dev.jsonl',
        ]
        for xpath in xlist:
            check(xpath, 0)

    _main()
